from typing import List

import numpy as np

from nn.layers import Linear, Tanh  # 线性层 + 激活函数
from nn.nn import NeuralNet  # 神经网络容器
from nn.optim import Adam  # 优化器
from nn.train import train  # 训练函数

# ================================
# 🔢 数据编码函数
# ================================


def binary_encode(x: int) -> List[int]:
    """
    将整数 x 编码为 10 位二进制向量(最低有效位在前)
    例如: x = 5 → 二进制 '101' → [1, 0, 1, 0, 0, 0, 0, 0, 0, 0]
    用途: 将数字转为神经网络输入(10维特征)
    """
    return [x >> i & 1 for i in range(10)]


def fizz_buzz_encode(x: int) -> List[int]:
    """
    将整数 x 编码为 4 维 one-hot 标签：
        [1, 0, 0, 0] → 普通数字(如 1, 2, 4...)
        [0, 1, 0, 0] → 能被 3 整除("fizz")
        [0, 0, 1, 0] → 能被 5 整除("buzz")
        [0, 0, 0, 1] → 能被 15 整除("fizzbuzz")
    """
    if x % 15 == 0:
        return [0, 0, 0, 1]  # fizzbuzz
    elif x % 5 == 0:
        return [0, 0, 1, 0]  # buzz
    elif x % 3 == 0:
        return [0, 1, 0, 0]  # fizz
    else:
        return [1, 0, 0, 0]  # 数字本身


# ================================
# 📊 构建训练数据集(101 ~ 1023)
# ================================

# 输入：每个数字 → 10 维二进制向量
inputs = np.array([binary_encode(x) for x in range(101, 1024)])
# inputs.shape = (923, 10)

# 标签：每个数字 → 4 维 one-hot 向量
targets = np.array([fizz_buzz_encode(x) for x in range(101, 1024)])
# targets.shape = (923, 4)


# ================================
# 🧠 构建神经网络模型
# ================================

net = NeuralNet(
    [
        # 第一层：线性变换，10 → 50
        # 输入形状: (batch_size, 10)
        # 输出形状: (batch_size, 50)
        Linear(input_size=10, output_size=50),
        # 激活函数：tanh，非线性
        # 输入/输出形状: (batch_size, 50)
        Tanh(),
        # 第二层：线性变换，50 → 4(4 类输出)
        # 输入形状: (batch_size, 50)
        # 输出形状: (batch_size, 4) —— logits(未归一化)
        Linear(input_size=50, output_size=4),
    ]
)

# 模型结构总结：
# (10) → Linear → (50) → Tanh → (50) → Linear → (4)


# ================================
# 🏃‍♂️ 开始训练
# ================================

# 使用 Adam 优化器 + 交叉熵损失(通常在 train 内部自动使用)
# train 函数会：
#   - 自动划分 batch
#   - 前向传播
#   - 计算 CrossEntropyLoss
#   - 反向传播
#   - 更新参数
train(
    net,
    inputs,  # (923, 10)
    targets,  # (923, 4) —— 注意：CrossEntropyLoss 通常期望整数标签，
    optimizer=Adam(),  # 自适应学习率优化器
)

# ================================
# 🧪 测试模型(1 ~ 100)
# ================================

print("\n" + "=" * 50)
print("🧪 FizzBuzz 测试结果(1 ~ 100)")
print("=" * 50)

num_correct = 0
labels = ["number", "fizz", "buzz", "fizzbuzz"]  # 对应类别索引 0~3

for x in range(1, 101):
    # 1. 编码输入数字为 10 维二进制
    input_vec = binary_encode(x)  # List[int], len=10

    # 2. 模型预测(前向传播)
    predicted = net(input_vec)  # 输出形状: (4,) —— logits 或概率

    # 3. 取最大值索引作为预测类别
    predicted_idx = np.argmax(predicted)  # 0, 1, 2, 或 3

    # 4. 获取真实类别索引
    actual_idx = np.argmax(fizz_buzz_encode(x))

    # 5. 判断是否正确
    if predicted_idx == actual_idx:
        num_correct += 1
        status = "✅"
    else:
        status = "❌"

    # 6. 打印结果
    print(
        f"{status} {x:2d}: 预测={labels[predicted_idx]:8s} | 真实={labels[actual_idx]}"
    )

# 打印最终准确率
print("\n" + "=" * 50)
print(f"🎯 正确数: {num_correct} / 100")
print(f"📈 准确率: {num_correct}%")
print("=" * 50)
